#include "genome.hh"
#include "constants.hh"
#include "mutations.hh"
#include <iostream>
#include <string>
#include <algorithm>
#include <sstream>
using namespace std;
using namespace evolution;

int Genome::nodes() const {
    return node_genes.size();
}

void
Genome::check_innov(int innov) {
    if (min_innov == -1 || min_innov >= innov) min_innov = innov;
    if (max_innov == -1 || max_innov < innov) max_innov = innov;
}


shared_ptr<Genome>
Genome::random_genome(int initial_nodes, int input_nodes) {
    srand(time(0));

    shared_ptr<Genome> genome = make_shared<Genome>(input_nodes);

    for (int i = 0; i < initial_nodes; ++i) {
        genome->add_node(NodeGene{i, Gene::random_weight(), NodeGene::random_type()});
    }

    
    for (int src = 0; src < initial_nodes; ++src) {
        for (int dst = src < input_nodes ? input_nodes : src+1;
                 dst < initial_nodes; ++dst) {
            if (with_probability(0.8f)) {
                const int innov = ConnectionGene::new_innovation_number();
                genome->check_innov(innov);
                genome->add_connection({innov,
                       src, dst, Gene::random_weight(), true});
            }
        }
    }

    return genome;
}

shared_ptr<Genome>
Genome::empty_genome(int initial_nodes, int input_nodes) {
    shared_ptr<Genome> genome = make_shared<Genome>(input_nodes);

    for (int i = 0; i < initial_nodes; ++i) {
        genome->add_node(NodeGene{i, 1.0f, NodeType::LINEAR});
    }

    return genome;
}

shared_ptr<Genome>
Genome::clone_genome(shared_ptr<Genome> const& other) {
    shared_ptr<Genome> genome = copy_genome(other);
    genome->mutate();
    return genome;
}

shared_ptr<Genome>
Genome::copy_genome(shared_ptr<Genome> const& other) {
    shared_ptr<Genome> genome = make_shared<Genome>(other->input_nodes);
    
    genome->max_innov = other->max_innov;
    genome->min_innov = other->min_innov;
    genome->node_genes = other->node_genes;
    genome->connection_genes = other->connection_genes;

    return genome;
}




void
Genome::mutate_add_node() {
    if (with_probability(MUTATION_ADD_NODE_PROBABILITY)) {
        if (connection_genes.size() == 0) {
            mutate_add_connection();
            return;
        }
        auto& connection = connection_genes[generator() % connection_genes.size()];

        const int new_id = node_genes.size();
        connection.set_enabled(false);

        list<NodeGene>::iterator pos = find_node(connection.get_source());
        if (pos != node_genes.end()) ++pos;
        if (distance(node_genes.begin(), pos) < input_nodes) {
            advance(pos, input_nodes - distance(node_genes.begin(), pos) + 1);
        }
        node_genes.insert(pos, {new_id, 1.0f, NodeType::LINEAR});

        const int a_innov = ConnectionGene::new_innovation_number();
        const int b_innov = ConnectionGene::new_innovation_number();
        ConnectionGene a {a_innov, connection.get_source(), new_id,
                            1.0f, true};
        ConnectionGene b {b_innov, new_id, connection.get_target(),
                            connection.get_weight(), true};
        check_innov(a_innov);
        check_innov(b_innov);

        add_connection(a);
        add_connection(b);
    }
}

void
Genome::mutate_add_connection() {
    if (with_probability(MUTATION_ADD_CONNECTION_PROBABILITY)) {
        list<NodeGene>::iterator src = node_genes.begin();
        const int src_offset = generator() % (node_genes.size() - 2);
        advance(src, src_offset);

        list<NodeGene>::iterator dst = src;
        int dst_offset = 1 + generator() % (node_genes.size() - src_offset - 1);
        advance(dst, dst_offset);

        for (auto& c: connection_genes) {
            if (c.get_source() == src->get_id() && c.get_target() == dst->get_id()) {
                return; //already exists!
            }
        }

        const int innov = ConnectionGene::new_innovation_number();
        check_innov(innov);
        add_connection({innov, src->get_id(), dst->get_id(),
                            Gene::random_weight(), true});
    }
}

void
Genome::mutate_weights() {
    for (auto& node: node_genes) node.mutate();
    for (auto& connection: connection_genes) connection.mutate();
}

void
Genome::mutate() {
    mutate_add_node();
    mutate_add_connection();
    mutate_weights();
}

list<NodeGene>::iterator 
Genome::find_node (int id) {
    return find_if (node_genes.begin(), node_genes.end(),
        [=](NodeGene const& node) {
            return node.get_id() == id;
        });
}








namespace evolution {
    ostream& operator << (ostream& out, Genome const& genome) {
        
        out << "Genome\n";

        for (auto node: genome.get_node_genes()) {
            out << "Node\n";
            out << node.get_id() << ' ' << node.get_weight() << ' '
                << static_cast<int>(node.get_type()) << '\n';
        }

        for (auto conn: genome.get_connection_genes()) {
            out << "Conn\n";
            out << conn.get_enabled() << ' ' << conn.get_source() << ' ' << conn.get_target() << ' '
                << conn.get_weight() << ' ' << conn.get_innovation_number() << '\n';
        }

        out << "End\n\n";

        return out;
    }

    istream& operator >> (istream& in, Genome& genome) {

        string line;
        do {
            getline(in, line);
            if (line == "End") return in;
            else if (line == "Genome") continue;   

            if (line == "Node") {
                getline(in, line);
                istringstream sin{line};
        
                int id; float bias; int type;
                sin >> id >> bias >> type;
                genome.node_genes.push_back({id, bias, static_cast<NodeType>(type)});

            } else if (line == "Conn") {
                getline(in, line);
                istringstream sin{line};

                int innov; float weight; int src; int dst; bool enabled;
                sin >> enabled >> src >> dst >> weight >> innov;

                genome.check_innov(innov);
                genome.connection_genes.push_back({innov, src, dst, weight, enabled});
            }
        } while (in);

        return in;
    }
}





